package com.ayu.shortlink.project.util;

import com.ayu.shortlink.project.common.cache.CaffeineCacheFactory;
import com.ayu.shortlink.project.common.cache.CommonCacheService;
import com.ayu.shortlink.project.common.enums.BucketTokenKeyEnum;
import jodd.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.CommandLineRunner;
import org.springframework.stereotype.Component;

import java.util.concurrent.Executors;

/**
 * 令牌桶限流器组件
 */
@Slf4j
@Component
public class CaffeineTokenBucketRateLimiter implements CommandLineRunner {
    // 令牌桶的容量
    private final int bucketCapacity = 1000;
    // 令牌生成的时间间隔（毫秒）
    private final long refillInterval = 2;
    // 每次生成的令牌数
    private final int refillTokens = 500;
    // 令牌桶
    private final CommonCacheService<String, Integer> tokenBucketCache = CaffeineCacheFactory.getTokenBucketCache();

    /**
     * 组件启动调用refillTokens()定时更新令牌
     * @param args args
     * @throws Exception 异常
     */
    @Override
    public void run(String... args) throws Exception {
        refillTokens();
    }

    /**
     * 尝试获取一个令牌
     * @param key key
     * @return true - 获取成功 false - 获取失败
     */
    public boolean tryAcquire(String key){
        Integer result = tokenBucketCache.getMap().compute(key, (k, v) -> {
            if (v == null || v < 1) {
                return -1;
            }
            return v - 1;
        });
        return result >= 0;
    }

    /**
     * 尝试获取n个令牌
     * @param key key
     * @param tokenNumber 令牌数量
     * @return true - 获取成功 false - 获取失败
     */
    public boolean tryAcquire(String key,Integer tokenNumber){
        Integer result = tokenBucketCache.getMap().compute(key, (k, v) -> {
            if (v == null || v < tokenNumber) {
                return -1;
            }
            return v - tokenNumber;
        });
        return result >= 0;
    }

    /**
     * 定时向令牌桶添加令牌
     */
    private void refillTokens(){
        Executors.newSingleThreadExecutor(ThreadFactoryBuilder
                        .create()
                        .setDaemon(true)
                        .setUncaughtExceptionHandler(new Thread.UncaughtExceptionHandler() {
                            @Override
                            public void uncaughtException(Thread t, Throwable e) {
                                // 报警
                                log.error("定时向令牌桶添加令牌出现异常");
                                e.printStackTrace();
                            }
                        })
                        .setNameFormat("thread-refillTokens")
                        .get())
                .execute(() -> {
                    while (true){
                        try {
                            // 进行令牌添加操作
                            for (BucketTokenKeyEnum x : BucketTokenKeyEnum.values()){
                                int oldTokens = tokenBucketCache.get(x.getKey()) == null ? 0 : tokenBucketCache.get(x.getKey());
                                int newTokens = Math.min(bucketCapacity, oldTokens + refillTokens);
                                tokenBucketCache.put(x.getKey(), newTokens);
                            }
                            Thread.sleep(refillInterval);
                        } catch (InterruptedException e) {
                            throw new RuntimeException(e);
                        }
                    }
        });
    }
}
